package com.stone.inventory.sharding.infra.algorithm;

import io.shardingsphere.api.algorithm.sharding.ListShardingValue;
import io.shardingsphere.api.algorithm.sharding.ShardingValue;
import io.shardingsphere.api.algorithm.sharding.complex.ComplexKeysShardingAlgorithm;
import lombok.Data;
import org.apache.commons.collections4.CollectionUtils;

import java.util.Collection;
import java.util.LinkedHashSet;

/**
 * sharding-jdbc默认的复合key分片算法
 *
 * @author Mr_wenpan@163.com
 */
@Data
public class DefaultComplexKeysShardingAlgorithm implements ComplexKeysShardingAlgorithm {

    /**
     * 主键列名称
     */
    private String idColumnName;

    /**
     * 分片列名称
     */
    private String shardingColumnName;

    /**
     * 分片算法
     */
    private ShardingAlgorithm shardingAlgorithm;

    public DefaultComplexKeysShardingAlgorithm(String idColumnName, String shardingColumnName, ShardingAlgorithm shardingAlgorithm) {
        this.idColumnName = idColumnName;
        this.shardingColumnName = shardingColumnName;
        this.shardingAlgorithm = shardingAlgorithm;
    }

    @Override
    public Collection<String> doSharding(Collection<String> availableTargetNames, Collection<ShardingValue> shardingValues) {

        if (CollectionUtils.isEmpty(shardingValues)) {
            return availableTargetNames;
        }
        Collection<String> result = new LinkedHashSet<>();
        Collection<Object> values = getShardingValues(shardingValues, idColumnName);
        if (CollectionUtils.isNotEmpty(values)) {
            values.forEach(id -> result.add(shardingAlgorithm.getTargetTableByShardingId((Long) id)));
            return result;
        }
        values = getShardingValues(shardingValues, shardingColumnName);
        if (CollectionUtils.isNotEmpty(values)) {

            values.forEach(shardingValue -> result.add(shardingAlgorithm.getTargetTableByShardingValue(shardingValue)));
            return result;
        }
        throw new RuntimeException("Error in Piecewise Algorithms: sharding data is ambiguity");
    }

    /**
     * 获取分片字段值
     */
    private Collection<Object> getShardingValues(Collection<ShardingValue> shardingValues, String column) {

        for (ShardingValue shardingValue : shardingValues) {
            if (!shardingValue.getColumnName().equals(column)) {
                continue;
            }
            if (shardingValue instanceof ListShardingValue) {
                return ((ListShardingValue) shardingValue).getValues();
            }
        }
        return null;
    }
}
